import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import griddata

# 生成二维数据集
np.random.seed(0)
x = np.random.rand(100) * 10
y = np.random.rand(100) * 10
z = np.sin(x/2) * np.cos(y/2)

# 添加一些NaN值模拟缺失数据
mask = np.random.choice([True, False], size=z.shape, p=[0.1, 0.9])
z[mask] = np.nan

# 定义插值网格
xi = np.linspace(0, 10, 100)
yi = np.linspace(0, 10, 100)
xi, yi = np.meshgrid(xi, yi)

# 使用最近邻插值填补缺失值
zi_nn = griddata((x[~mask], y[~mask]), z[~mask], (xi, yi), method='nearest')

# 创建图形窗口
fig, axs = plt.subplots(1, 3, figsize=(18, 6))

# 1. 原始数据分布图（包含NaN）
sc1 = axs[0].scatter(x, y, c=z, cmap='jet', s=50)
axs[0].set_title('Original Data with NaN')
axs[0].set_xlabel('X-axis')
axs[0].set_ylabel('Y-axis')
fig.colorbar(sc1, ax=axs[0])

# 2. 最近邻插值后数据分布图
sc2 = axs[1].contourf(xi, yi, zi_nn, cmap='rainbow', levels=100)
axs[1].set_title('Data after Nearest Neighbor Interpolation')
axs[1].set_xlabel('X-axis')
axs[1].set_ylabel('Y-axis')
fig.colorbar(sc2, ax=axs[1])

# 3. 插值前后差异图
zi_original = griddata((x, y), z, (xi, yi), method='nearest', fill_value=np.nan)
diff = np.abs(zi_nn - zi_original)
sc3 = axs[2].contourf(xi, yi, diff, cmap='plasma', levels=100)
axs[2].set_title('Difference between Original and Interpolated Data')
axs[2].set_xlabel('X-axis')
axs[2].set_ylabel('Y-axis')
fig.colorbar(sc3, ax=axs[2])

# 调整图像布局
plt.tight_layout()
plt.show()
